import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Union

import torch
from torch.nn import MSELoss, CrossEntropyLoss
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader
from torchvision import transforms

from zsl_ma.dataset_utils.CustomImageDataset import ImageClassificationDataset
from zsl_ma.models.VAE import DRCAE
from zsl_ma.tools.tool import create_csv, append_metrics_to_csv, get_device, setup_logger, setup_save_dirs
from zsl_ma.tools.train_val_until import train_cae_one_epoch

logger = logging.getLogger('zsl_train')


def train_cae(configs):
    device = configs.device
    transform = configs.transform
    save_dir = configs.save_dir
    results_file = os.path.join(save_dir, 'cae_metrics.csv')
    metrics = ['epoch', 'train_loss', 'val_loss', 'lr']
    create_csv(metrics, results_file)

    train_dataset = ImageClassificationDataset(configs.data_dir, transform=transform, train_class=configs.train_class)
    val_dataset = ImageClassificationDataset(configs.data_dir, transform=transform, train_class=configs.train_class,
                                             mode='val')
    train_loader = DataLoader(train_dataset, batch_size=configs.batch_size, shuffle=True,
                              num_workers=configs.num_workers)
    val_loader = DataLoader(val_dataset, batch_size=configs.batch_size, shuffle=False, num_workers=configs.num_workers)

    model = DRCAE().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=configs.lr)
    lr_scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.1, patience=5, min_lr=1e-9)
    criterion = MSELoss()

    best = 1e8
    num_epochs = configs.epochs

    patience = configs.patience  # 从外部参数获取耐心值
    early_stop_counter = 0  # 早停计数器
    best_epoch = 0  # 最佳模型的epoch
    for epoch in range(num_epochs):
        training_lr = lr_scheduler.get_last_lr()[0]

        train_loss, val_loss = train_cae_one_epoch(model, train_loader, val_loader, device, optimizer, criterion, epoch)
        lr_scheduler.step(val_loss)

        logger.info(f'the epoch {epoch + 1} train loss is {train_loss:.6f}, val loss is {val_loss:.6f}')

        metric = {'epoch': epoch, 'train_loss': train_loss, 'val_loss': val_loss, 'lr': training_lr}
        append_metrics_to_csv(metric, results_file)

        current_score = val_loss
        if current_score < best:
            best = current_score
            best_epoch = epoch
            early_stop_counter = 0  # 重置早停计数器
            model.save_encoder(os.path.join(save_dir, 'checkpoints', 'encoder.pth'))
            logger.info(f'Best model saved at epoch {epoch + 1} with loss: {best:.6f}')
        elif patience > 0 and training_lr < (configs.lr * 0.0001):  # 仅当patience>0时执行早停计数
            early_stop_counter += 1
            logger.info(f'Early stopping counter: {early_stop_counter}/{patience}')

            # 如果早停计数器达到耐心值，停止训练
            if early_stop_counter >= patience:
                logger.info(f'Early stopping triggered after {epoch + 1} epochs')
                break
    logger.info(f'Best model was saved at epoch {best_epoch + 1} with F1-score: {best:.6f}')


def get_cae_args(args=None):
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--data_dir', type=str,
                        default=r'D:\Code\2-ZSL\0-data\CWRU\dataset')
    parser.add_argument('--save_dir', type=str, default=r'D:\Code\2-ZSL\1-output\特征解耦结果\exp-3')
    parser.add_argument('--train_class', type=str,
                        default=r'D:\Code\2-ZSL\0-data\CWRU\dataset/seen_classes.txt')
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--batch_size', default=100, type=int)
    parser.add_argument('--lr', default=1e-1, type=float)
    parser.add_argument('--weight_decay', default=1e-5, type=float)
    parser.add_argument('--patience', default=10, type=int)
    parser.add_argument('--num_workers', default=0, type=int)

    return parser.parse_args(args if args else [])


@dataclass
class TrainCaeConfig:
    """CAE训练配置参数类"""
    device: torch.device = field(
        default_factory=get_device,  # 用default_factory延迟初始化，避免模块导入时执行
        metadata={"desc": "训练使用的设备（CPU/GPU）"}
    )
    transform: transforms.Compose = field(
        default_factory=lambda: transforms.Compose([  # 用default_factory创建Compose实例
            transforms.Resize((64, 64)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]),
        metadata={"desc": "图像预处理变换管道"}
    )
    # 数据路径参数
    data_dir: Union[str, Path] = r'D:\Code\2-ZSL\0-data\CWRU\dataset'
    save_dir: Union[str, Path] = r'D:\Code\2-ZSL\1-output\特征解耦结果\exp-3'
    train_class: Union[str, Path] = r'D:\Code\2-ZSL\0-data\CWRU\dataset/seen_classes.txt'

    # 训练超参数
    epochs: int = 1500
    batch_size: int = 100
    lr: float = 1e-1
    weight_decay: float = 1e-5
    patience: int = 10
    num_workers: int = 0


if __name__ == '__main__':
    opts = TrainCaeConfig(data_dir=r'/data/coding/dataset',
                          save_dir=r'/data/coding/output/H03/exp-1',
                          train_class=r'/data/coding/dataset/seen_classes.txt',
                          lr=0.01,
                          batch_size=128,
                          num_workers=int(32))
    opts.save_dir = setup_save_dirs(opts.save_dir, None)
    logger = setup_logger(opts.save_dir)
    # 打印配置并启动训练（此时会触发设备信息日志）
    logger.info("=" * 60)
    logger.info("【单独运行】分类器训练流程")
    logger.info(f"配置信息:{vars(opts)}")
    logger.info("=" * 60)
    train_cae(opts)
